import socket, struct, time, threading


class HeartBeat():
    def __init__(self):
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.send_addr = ('192.168.1.120', 43893)
        self.heart_task = threading.Thread(target=self.send_heartbeat)
        self.heart_task.start()

    def send_heartbeat(self):
        while True:
            data = struct.pack("<3i", 0x21040001, 0, 0)
            self.udp_socket.sendto(data, self.send_addr)
            time.sleep(0.5)  #sending command frequency not lower than 2HZ


class RobotState:
    def __init__(self):
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_socket.bind(('192.168.1.103', 43897))  # 使用不同端口接收状态信息
        
        # 基础状态
        self.power = False  # 是否开机
        self.is_standing = False  # 是否站立
        self.current_mode = None  # 'spot', 'move', 'autonomous', 'manual'
        
        # 机器人状态信息 (根据PDF第12页结构体)
        self.robot_mode = 0  # 机器人模式
        self.gait_type = 0  # 步态类型
        self.battery_level = 0  # 电池电量
        self.battery_current = 0  # 电池电流
        self.battery_temperature = 0  # 电池温度
        self.battery_voltage = 0  # 电池电压
        
        # 位置和姿态
        self.position = [0.0, 0.0, 0.0]  # x, y, z位置
        self.velocity = [0.0, 0.0, 0.0]  # x, y, z速度
        self.angular_velocity = [0.0, 0.0, 0.0]  # roll, pitch, yaw角速度
        self.attitude = [0.0, 0.0, 0.0]  # roll, pitch, yaw姿态角
        
        # 足端状态
        self.foot_force = [0, 0, 0, 0]  # 四足着地状态
        self.foot_force_est = [0.0, 0.0, 0.0, 0.0]  # 四足估计压力
        self.foot_position = [[0.0, 0.0, 0.0] for _ in range(4)]  # 四足位置
        self.foot_speed = [[0.0, 0.0, 0.0] for _ in range(4)]  # 四足速度
        
        # 电机状态
        self.motor_temperature = [0] * 12  # 12个电机温度
        self.motor_speed = [0] * 12  # 12个电机速度
        self.motor_position = [0] * 12  # 12个电机位置
        self.motor_current = [0] * 12  # 12个电机电流
        self.motor_voltage = [0] * 12  # 12个电机电压
        
        # 错误状态
        self.error_code = 0  # 错误代码
        self.warning_code = 0  # 警告代码
        
        # 启动状态监控线程
        self.monitor_thread = threading.Thread(target=self.receive_state, daemon=True)
        self.monitor_thread.start()

    def receive_state(self):
        """接收机器狗状态信息"""
        while True:
            try:
                data, addr = self.udp_socket.recvfrom(1024)
                self._parse_state_data(data)
            except Exception as e:
                print(f"状态接收错误: {e}")
            time.sleep(0.01)

    def _parse_state_data(self, data):
        """解析状态数据"""
        try:
            # 解包数据头
            cmd_id, = struct.unpack("<I", data[0:4])
            
            if cmd_id == 0x21040001:  # 状态数据命令ID
                offset = 4  # 跳过命令ID
                
                # 解析基础状态
                (self.robot_mode, self.gait_type, self.battery_level, 
                 self.battery_current, self.battery_temperature, self.battery_voltage) = struct.unpack("<6i", data[offset:offset+24])
                offset += 24
                
                # 解析位置和姿态
                self.position = list(struct.unpack("<3f", data[offset:offset+12]))
                offset += 12
                self.velocity = list(struct.unpack("<3f", data[offset:offset+12]))
                offset += 12
                self.angular_velocity = list(struct.unpack("<3f", data[offset:offset+12]))
                offset += 12
                self.attitude = list(struct.unpack("<3f", data[offset:offset+12]))
                offset += 12
                
                # 解析足端状态
                self.foot_force = list(struct.unpack("<4i", data[offset:offset+16]))
                offset += 16
                self.foot_force_est = list(struct.unpack("<4f", data[offset:offset+16]))
                offset += 16
                
                # 解析四足位置和速度
                for i in range(4):
                    self.foot_position[i] = list(struct.unpack("<3f", data[offset:offset+12]))
                    offset += 12
                for i in range(4):
                    self.foot_speed[i] = list(struct.unpack("<3f", data[offset:offset+12]))
                    offset += 12
                
                # 解析电机状态
                self.motor_temperature = list(struct.unpack("<12i", data[offset:offset+48]))
                offset += 48
                self.motor_speed = list(struct.unpack("<12f", data[offset:offset+48]))
                offset += 48
                self.motor_position = list(struct.unpack("<12f", data[offset:offset+48]))
                offset += 48
                self.motor_current = list(struct.unpack("<12f", data[offset:offset+48]))
                offset += 48
                self.motor_voltage = list(struct.unpack("<12f", data[offset:offset+48]))
                offset += 48
                
                # 解析错误和警告代码
                self.error_code, self.warning_code = struct.unpack("<2I", data[offset:offset+8])
                
                # 更新派生状态
                self.power = True  # 如果能收到数据就说明开机了
                self.is_standing = self.robot_mode in [1, 2]  # 假设模式1和2表示站立状态
                self.current_mode = {
                    0: None,
                    1: 'spot',
                    2: 'move',
                    3: 'autonomous',
                    4: 'manual'
                }.get(self.robot_mode, None)
                
                # 更新运动状态
                self.roll, self.pitch, self.yaw = self.attitude
                self.speed = (self.velocity[0]**2 + self.velocity[1]**2)**0.5  # 计算水平速度
                self.is_moving = self.speed > 0.01  # 速度大于阈值认为在运动
                
        except Exception as e:
            print(f"状态解析错误: {e}")

    # 状态检查方法
    def check_power(self):
        """检查电源状态"""
        if not self.power:
            return False, "机器狗未开机"
        return True, "电源正常"

    def check_battery(self, min_level=10):
        """检查电池电量"""
        if self.battery_level < min_level:
            return False, f"电池电量过低 ({self.battery_level}%)"
        return True, "电池电量正常"

    def check_emergency(self):
        """检查急停状态"""
        if self.is_emergency:
            return False, "机器狗处于急停状态"
        return True, "未处于急停状态"

    def check_standing(self):
        """检查站立状态"""
        if not self.is_standing:
            return False, "机器狗必须处于站立状态"
        return True, "站立状态正常"

    def check_mode(self, required_mode):
        """检查运动模式"""
        if self.current_mode != required_mode:
            return False, f"需要在{required_mode}模式下执行"
        return True, "运动模式正常"

    def check_motion(self, action_name):
        """检查运动状态"""
        if self.is_moving:
            return False, f"执行{action_name}时机器狗必须静止"
        return True, "运动状态正常"

    def check_posture(self, max_angle=30):
        """检查姿态角度"""
        if abs(self.roll) > max_angle or abs(self.pitch) > max_angle:
            return False, f"姿态角度过大 (roll:{self.roll}, pitch:{self.pitch})"
        return True, "姿态角度正常"

    def check_for_movement(self):
        """检查移动条件"""
        checks = [
            self.check_power(),
            self.check_battery(),
            self.check_emergency(),
            self.check_standing(),
            self.check_posture()
        ]
        for ok, msg in checks:
            if not ok:
                return False, msg
        return True, "可以移动"

    def check_for_backflip(self):
        """检查后空翻条件"""
        checks = [
            self.check_power(),
            self.check_battery(50),  # 后空翻需要更高的电量
            self.check_emergency(),
            self.check_standing(),
            self.check_motion("后空翻"),
            self.check_posture(20)  # 后空翻需要更严格的姿态要求
        ]
        for ok, msg in checks:
            if not ok:
                return False, msg
        return True, "可以执行后空翻"


class ControlRobot:
    def __init__(self, ip='192.168.1.120', port=43893):
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.send_addr = (ip, port)
        self.heartbeat = HeartBeat()
        self.state = RobotState()

    def send_data(self, code, value, type):
        ok, msg = self.state.check_power()
        if not ok:
            print(f"错误：{msg}")
            return False
        try:
            data = struct.pack("<3i", code, value, type)
            self.udp_socket.sendto(data, self.send_addr)
            return True
        except Exception as e:
            print(f"发送命令失败: {e}")
            return False

    # 基础控制命令
    def stand_up_or_lie_down(self):
        """切换站立/趴下状态"""
        ok, msg = self.state.check_power()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010202, 0, 0):
            print(f"执行{'站立' if not self.state.is_standing else '趴下'}命令")

    def soft_emergency_stop(self):
        """软急停"""
        if self.send_data(0x21020C0E, 0, 0):
            print("执行软急停")

    def zero_position(self):
        """回零位"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010C05, 0, 0):
            print("执行回零位")

    # 姿态调整命令（原地模式）
    def adjust_roll(self, value):
        """调整横滚角"""
        checks = [
            self.state.check_standing(),
            self.state.check_mode('spot')
        ]
        for ok, msg in checks:
            if not ok:
                print(f"错误：{msg}")
                return
        if self.send_data(0x21010131, value, 0):
            print(f"调整横滚角，值：{value}")

    def adjust_pitch(self, value):
        """调整俯仰角"""
        checks = [
            self.state.check_standing(),
            self.state.check_mode('spot')
        ]
        for ok, msg in checks:
            if not ok:
                print(f"错误：{msg}")
                return
        if self.send_data(0x21010130, value, 0):
            print(f"调整俯仰角，值：{value}")

    def adjust_height(self, value):
        """调整身体高度"""
        checks = [
            self.state.check_standing(),
            self.state.check_mode('spot')
        ]
        for ok, msg in checks:
            if not ok:
                print(f"错误：{msg}")
                return
        if self.send_data(0x21010102, value, 0):
            print(f"调整身体高度，值：{value}")

    def adjust_yaw(self, value):
        """调整偏航角"""
        checks = [
            self.state.check_standing(),
            self.state.check_mode('spot')
        ]
        for ok, msg in checks:
            if not ok:
                print(f"错误：{msg}")
                return
        if self.send_data(0x21010135, value, 0):
            print(f"调整偏航角，值：{value}")

    # 移动控制命令
    def move_side(self, value):
        """左右平移"""
        # ok, msg = self.state.check_for_movement()
        ok=True
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010131, value, 0):
            print(f"左右平移，速度值：{value}")

    def move_forward(self, value):
        """前后移动"""
        # ok, msg = self.state.check_for_movement()
        ok=True
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010130, value, 0):
            print(f"前后移动，速度值：{value}")

    def turn(self, value):
        """转向"""
        ok, msg = self.state.check_for_movement()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010135, value, 0):
            print(f"转向，速度值：{value}")

    # 模式切换命令
    def on_spot_mode(self):
        """切换到原地模式"""
        # ok, msg = self.state.check_for_movement()
        ok=True
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010D05, 0, 0):
            self.state.current_mode = 'spot'
            print("切换到原地模式")

    def move_mode(self):
        """切换到移动模式"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010D06, 0, 0):
            self.state.current_mode = 'move'
            print("切换到移动模式")

    # 步态控制命令
    def flat_low_speed_gait(self):
        """平地低速步态"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010300, 0, 0):
            self.state.current_gait = 'flat_low'
            print("切换到平地低速步态")

    def flat_medium_speed_gait(self):
        """平地中速步态"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010307, 0, 0):
            self.state.current_gait = 'flat_medium'
            print("切换到平地中速步态")

    def flat_high_speed_gait(self):
        """平地高速步态"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010303, 0, 0):
            self.state.current_gait = 'flat_high'
            print("切换到平地高速步态")

    # 特殊动作命令
    def twist_body(self):
        """扭身体"""
        ok, msg = self.state.check_for_movement()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010204, 0, 0):
            print("执行扭身体动作")

    def flip(self):
        """翻身"""
        ok, msg = self.state.check_for_backflip()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010205, 0, 0):
            print("执行翻身动作")

    def space_step(self):
        """太空步"""
        ok, msg = self.state.check_for_movement()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x2101030C, 0, 0):
            print("执行太空步")

    def backflip(self):
        """后空翻"""
        ok, msg = self.state.check_for_backflip()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010502, 0, 0):
            print("执行后空翻")

    def greet(self):
        """打招呼"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010507, 0, 0):
            print("执行打招呼动作")

    # 系统控制命令
    def autonomous_mode(self):
        """切换到自主模式"""
        # ok, msg = self.state.check_for_movement()
        ok=True
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010C03, 0, 0):
            self.state.current_mode = 'autonomous'
            print("切换到自主模式")

    def manual_mode(self):
        """切换到手动模式"""
        ok, msg = self.state.check_standing()
        if not ok:
            print(f"错误：{msg}")
            return
        if self.send_data(0x21010C02, 0, 0):
            self.state.current_mode = 'manual'
            print("切换到手动模式")

    # AI功能控制
    def close_all_ai_options(self):
        """关闭所有AI选项"""
        if self.send_data(0x21012109, 0x00, 0):
            print("关闭所有AI功能")

    def enable_obstacle_stop(self):
        """开启停障"""
        if self.send_data(0x21012109, 0x20, 0):
            print("开启停障功能")

    def enable_follow(self):
        """开启跟随"""
        if self.send_data(0x21012109, 0xC0, 0):
            print("开启跟随功能")
# def main():
#     controller = ControlRobot()
#     time.sleep(1)  # 等待初始化完成
#     while True:
#         print("power:", controller.state.power)
# if __name__ == "__main__":
#     main()